package cn.ycc1.lianzi.controller;

import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import io.github.pigmesh.ai.deepseek.core.DeepSeekClient;
import io.github.pigmesh.ai.deepseek.core.EmbeddingClient;
import io.github.pigmesh.ai.deepseek.core.chat.ChatCompletionRequest;
import io.github.pigmesh.ai.deepseek.core.chat.ChatCompletionResponse;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.CreateDatabaseParam;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.DropCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.SearchResp;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

/**
 * DeepSeek RAG 练习
 * https://github.com/pig-mesh/deepseek4j
 * 运行 Deepseek 调试工具 sse.html
 *
 * Milvus 官方文档：https://milvus.io/docs/
 * https://milvus.io/docs/drop-collection.md
 *
 * client = MilvusClient(uri="http://localhost:19530", token="root:Milvus")
 * Milvus linux docker安装
 * curl -sfL https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh -o standalone_embed.sh
 * bash standalone_embed.sh start
 * You can access Milvus WebUI at http://127.0.0.1:9091/webui/ to learn more about the your Milvus instance. For details, refer to Milvus WebUI.
 * @author ycc
 * @date 2025/2/20
 */
@RestController
public class DeepSeekController {
    @Autowired
    private DeepSeekClient deepSeekClient;

    @Autowired
    EmbeddingClient embeddingClient;

    @Autowired
    private Optional<EmbeddingClient> embeddingClientOptional;

    private static final String CLUSTER_ENDPOINT = "http://"+ System.getenv("DEEPSEEK_MILVUS_HOST") +":19530";
    private static final String TOKEN = "root:Milvus";

    // sse流式返回
    @GetMapping(value = "/deepseek/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatCompletionResponse> chat(String prompt) {
        return deepSeekClient.chatFluxCompletion(prompt);
    }

    @GetMapping(value = "/deepseek/chat2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatCompletionResponse> chat2(String prompt) {
        ChatCompletionRequest request = ChatCompletionRequest.builder()
                // 模型选择，支持 DEEPSEEK_CHAT、DEEPSEEK_REASONER 等
                // .model(ChatCompletionModel.DEEPSEEK_CHAT)
                // 添加用户消息
                .addUserMessage(prompt)
                // 添加助手消息，用于多轮对话
                .addAssistantMessage("上轮结果")
                // 添加系统消息，用于设置角色和行为
                .addSystemMessage("你是一个专业的助手")
                // 设置最大生成 token 数，默认 2048
                .maxTokens(1000)
                // 设置响应格式，支持 JSON 结构化输出
                //.responseFormat()
                .tools() // function calling
                .build();

        return deepSeekClient.chatFluxCompletion(request);
    }

    /**
     * 创建数据库到 Milvus
     */
    @GetMapping("/deepseek/create-database")
    public void createDatabase() {
        // 1. Connect to Milvus server
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withUri(CLUSTER_ENDPOINT)
                .withToken(TOKEN)
                .build();

        MilvusServiceClient client = new MilvusServiceClient(connectParam);

        // 3. Create a new database
        CreateDatabaseParam createDatabaseParam = CreateDatabaseParam.newBuilder()
                .withDatabaseName("my_database")
                .build();

        R<RpcStatus> response = client.createDatabase(createDatabaseParam);
        System.out.println(response);
    }

    /**
     * 创建集合到 Milvus
     */
    @GetMapping("/deepseek/create-collection")
    public void createCollection() {
        ConnectConfig connectConfig = ConnectConfig.builder()
                .uri(CLUSTER_ENDPOINT) // 1.2 获取的 Milvus 链接端点
                .token(TOKEN)  // 1.2 获取的 Milvus 链接信息
                .build();
        MilvusClientV2 client = new MilvusClientV2(connectConfig);

        // 3.1 Create schema
        CreateCollectionReq.CollectionSchema schema = client.createSchema();

        // 3.2 Add fields to schema
        schema.addField(AddFieldReq.builder()
                .fieldName("my_id")
                .dataType(DataType.Int64)
                .isPrimaryKey(true)
                .autoID(true)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("my_vector")
                .dataType(DataType.FloatVector)
                .dimension(1024)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("my_varchar")
                .dataType(DataType.VarChar)
                .maxLength(2048)
                .build());

        // 3.3 Prepare index parameters
        IndexParam indexParamForIdField = IndexParam.builder()
                .fieldName("my_id")
                .indexType(IndexParam.IndexType.STL_SORT)
                .build();

        IndexParam indexParamForVectorField = IndexParam.builder()
                .fieldName("my_vector")
                .indexType(IndexParam.IndexType.AUTOINDEX)
                .metricType(IndexParam.MetricType.COSINE)
                .build();

        List<IndexParam> indexParams = new ArrayList<>();
        indexParams.add(indexParamForIdField);
        indexParams.add(indexParamForVectorField);

        // 3.4 Create a collection with schema and index parameters
        CreateCollectionReq customizedSetupReq1 = CreateCollectionReq.builder()
                .collectionName("deepseek4j_test")
                .collectionSchema(schema)
                .indexParams(indexParams)
                .build();

        client.createCollection(customizedSetupReq1);

        // 3.5 Get load state of the collection
        GetLoadStateReq customSetupLoadStateReq1 = GetLoadStateReq.builder()
                .collectionName("deepseek4j_test")
                .build();

        Boolean loaded = client.getLoadState(customSetupLoadStateReq1);
        System.out.println(loaded);

        // Output:
        // true
    }

    /**
     * 写入向量到 Milvus
     */
    @GetMapping("/deepseek/add-milvus")
    public void addMilvus() {
        ConnectConfig connectConfig = ConnectConfig.builder()
                .uri(CLUSTER_ENDPOINT) // 1.2 获取的 Milvus 链接端点
                .token(TOKEN)  // 1.2 获取的 Milvus 链接信息
                .build();
        MilvusClientV2 milvusClientV2 = new MilvusClientV2(connectConfig);

        // 这里以 2025最新的我司保密条例演示，可以换成你自己的
        // String law = FileUtil.readString("E:\\test\\law.txt", Charset.defaultCharset());
        //String[] lawSplits = StrUtil.split(law, 400);
        String[] lawSplits = {
                "《江南好》江南好，水润吴江城。日出同里游古镇，归来姚家赏太湖。雨期几多情？",
                "《莲子图》莲子露，二十四节气。循环往复时令经，生生不息气候脉。斗转星移图？",
                "《五行》池边生草木，大地络五气。天上布星宿，循环经土圭。我本有阴阳，五行循环之。阴阳不离散，自然调和之。",
                "《闲走》风吹大观湖，心中惬意生。愁云暂散去，静享湿地风。",
                "《游流溪》林荫露天日，烟雨沐流溪。曲径通幽路，处处听流笛。",
                "《白云山》白云山，清风拂亭晚。劳逸结合养习惯，积极进取生经验。生活入自然。"
        };

        List<JsonObject> data = new ArrayList<>();
        for (String lawSplit : lawSplits) {
            List<Float> floatList = embeddingClient.embed(lawSplit);

            JsonObject jsonObject = new JsonObject();

            // 将 List<Float> 转换为 JsonArray
            JsonArray jsonArray = new JsonArray();
            for (Float value : floatList) {
                jsonArray.add(value);
            }
            jsonObject.add("my_vector", jsonArray);
            jsonObject.addProperty("my_varchar", lawSplit);

            data.add(jsonObject);
        }

        InsertReq insertReq = InsertReq.builder()
                .collectionName("deepseek4j_test")
                .data(data)
                .build();

        milvusClientV2.insert(insertReq);
    }

    /**
     * 删除集合到 Milvus
     */
    @GetMapping("/deepseek/drop-collection")
    public void dropCollection() {
        ConnectConfig connectConfig = ConnectConfig.builder()
                .uri(CLUSTER_ENDPOINT) // 1.2 获取的 Milvus 链接端点
                .token(TOKEN)  // 1.2 获取的 Milvus 链接信息
                .build();
        MilvusClientV2 milvusClientV2 = new MilvusClientV2(connectConfig);

        DropCollectionReq dropCollectionReq = DropCollectionReq.builder()
                .collectionName("deepseek4j_test")
                .build();

        milvusClientV2.dropCollection(dropCollectionReq);
    }

    /**
     * 创建 RAG 接口
     * @param prompt
     * @return
     */
    @GetMapping(value = "/deepseek/chat-rag", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatCompletionResponse> chatRag(String prompt) {
        ConnectConfig connectConfig = ConnectConfig.builder()
        .uri(CLUSTER_ENDPOINT) // 1.2 获取的 Milvus 链接端点
        .token(TOKEN)  // 1.2 获取的 Milvus 链接信息
        .build();

        MilvusClientV2 milvusClientV2 = new MilvusClientV2(connectConfig);

        List<Float> floatList = embeddingClientOptional.get().embed(prompt);

        SearchReq searchReq = SearchReq.builder()
                .collectionName("deepseek4j_test")
                .data(Collections.singletonList(new FloatVec(floatList)))
                .outputFields(Collections.singletonList("my_varchar"))
                .topK(3)
                .build();

        SearchResp searchResp = milvusClientV2.search(searchReq);

        List<String> resultList = new ArrayList<>();
        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        for (List<SearchResp.SearchResult> results : searchResults) {
            System.out.println("TopK results:");
            for (SearchResp.SearchResult result : results) {
                resultList.add(result.getEntity().get("my_varchar").toString());
            }
        }

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                // 根据渠道模型名称动态修改这个参数
                //.model("deepseek-r1:8b")
                //.model("deepseek-r1")
                .addUserMessage(String.format("你要根据用户输入的问题：%s \n \n 参考如下内容： %s  \n\n 整理处理最终结果", prompt, resultList)).build();

        return deepSeekClient.chatFluxCompletion(request);
    }

}
